package A一周刷爆LeetCode.B基础提升.B并查集;

import java.util.HashMap;
import java.util.List;
import java.util.Stack;

/**
 * @author wei.zhao
 * @description: 并查集
 * @date 2022年07月01日 15:51
 */
public class Code02_UnionFind {

    public static class UnionFindSet<T> {

        public HashMap<T, T> fatherMap;
        public HashMap<T, Integer> rankMap;

        public UnionFindSet(List<T> list) {
            for (T t : list) {
                fatherMap.put(t, t);
                rankMap.put(t, 1);
            }
        }

        // 是否在同一个集合中
        public boolean isSameSet(T t1, T t2) {
            if (fatherMap.containsKey(t1) && fatherMap.containsKey(t2) && findHead(t1) == findHead(t2)) {
                return true;
            }
            return false;
        }

        // 合并操作
        public void union(T t1, T t2) {
            if (isSameSet(t1, t2)) {
                return;
            }
            T head1 = findHead(t1);
            T head2 = findHead(t2);
            T lhead = rankMap.get(head1) >= rankMap.get(head2) ? head2 : head1;
            T ghead = rankMap.get(head1) >= rankMap.get(head2) ? head1 : head2;
            fatherMap.put(lhead, ghead);
            rankMap.put(ghead, rankMap.get(ghead) + rankMap.get(lhead));
            rankMap.remove(lhead);
        }

        // 查找头节点
        public T findHead(T t) {
            Stack<T> stack = new Stack<>();
            while (fatherMap.get(t) != t) {
                stack.add(t);
                t = fatherMap.get(t);
            }
            while (!stack.isEmpty()) {
                fatherMap.put(stack.pop(), t);
            }
            return t;
        }

    }

}
